package onnxruntime

import (
	"errors"
	"fmt"
)

// #cgo CFLAGS: -O2 -g
//
// #include "onnxruntime_wrapper.h"
import "C"

// Initializes and returns a SessionOptions struct, used when setting options
// in new AdvancedSession instances. The caller must call the Destroy()
// function on the returned struct when it's no longer needed.
func NewSessionOptions() (*SessionOptions, error) {
	if !IsInitialized() {
		return nil, ErrorNotInitialized
	}
	var o *C.OrtSessionOptions
	status := C.CreateSessionOptions(&o)
	if status != nil {
		return nil, statusToError(status)
	}
	return &SessionOptions{
		o: o,
	}, nil
}

// Used to set options when creating an ONNXRuntime session. There is currently
// not a way to change options after the session is created, apart from
// destroying the session and creating a new one. This struct opaquely wraps a
// C OrtSessionOptions struct, which users must modify via function calls. (The
// OrtSessionOptions struct is opaque in the C API, too.)
//
// Users must instantiate this struct using the NewSessionOptions function.
// Instances must be destroyed by calling the Destroy() method after the
// options are no longer needed (after NewAdvancedSession(...) has returned).
type SessionOptions struct {
	o *C.OrtSessionOptions
}

func (o *SessionOptions) Destroy() error {
	if o.o == nil {
		return errors.New("the SessionOptions are not initialized")
	}
	C.ReleaseSessionOptions(o.o)
	o.o = nil
	return nil
}

// Sets the number of threads used to parallelize execution within onnxruntime
// graph nodes. A value of 0 uses the default number of threads.
func (o *SessionOptions) SetIntraOpNumThreads(n int) error {
	if n < 0 {
		return fmt.Errorf("number of threads must be at least 0, got %d", n)
	}
	status := C.SetIntraOpNumThreads(o.o, C.int(n))
	if status != nil {
		return statusToError(status)
	}
	return nil
}

// Sets the number of threads used to parallelize execution across separate
// onnxruntime graph nodes. A value of 0 uses the default number of threads.
func (o *SessionOptions) SetInterOpNumThreads(n int) error {
	if n < 0 {
		return fmt.Errorf("number of threads must be at least 0, got %d", n)
	}
	status := C.SetInterOpNumThreads(o.o, C.int(n))
	if status != nil {
		return statusToError(status)
	}
	return nil
}

// Enable/Disable the usage of the memory arena on CPU.
// Arena may pre-allocate memory for future usage.
func (o *SessionOptions) SetCpuMemArena(isEnabled bool) error {
	n := 0
	if isEnabled {
		n = 1
	}
	status := C.SetCpuMemArena(o.o, C.int(n))
	if status != nil {
		return statusToError(status)
	}
	return nil
}

// Enable/Disable the memory pattern optimization.
// If this is enabled memory is preallocated if all shapes are known.
func (o *SessionOptions) SetMemPattern(isEnabled bool) error {
	n := 0
	if isEnabled {
		n = 1
	}
	status := C.SetMemPattern(o.o, C.int(n))
	if status != nil {
		return statusToError(status)
	}
	return nil
}

// Takes a pointer to an initialized CUDAProviderOptions instance, and applies
// them to the session options. This is what you'll need to call if you want
// the session to use CUDA. Returns an error if your device (or onnxruntime
// library) does not support CUDA. The CUDAProviderOptions struct can be
// destroyed after this.
func (o *SessionOptions) AppendExecutionProviderCUDA(
	cudaOptions *CUDAProviderOptions) error {
	status := C.AppendExecutionProviderCUDAV2(o.o, cudaOptions.o)
	if status != nil {
		return statusToError(status)
	}
	return nil
}

// Takes an initialized TensorRTProviderOptions instance, and applies them to
// the session options. You'll need to call this if you want the session to use
// TensorRT. Returns an error if your device (or onnxruntime library version)
// does not support TensorRT. The TensorRTProviderOptions can be destroyed
// after this.
func (o *SessionOptions) AppendExecutionProviderTensorRT(
	tensorRTOptions *TensorRTProviderOptions) error {
	status := C.AppendExecutionProviderTensorRTV2(o.o, tensorRTOptions.o)
	if status != nil {
		return statusToError(status)
	}
	return nil
}

// Enables the CoreML backend for the given session options on supported
// platforms. Unlike the other AppendExecutionProvider* functions, this one
// only takes a bitfield of flags rather than an options object, though it
// wouldn't suprise me if onnxruntime deprecated this API in the future as it
// did with the others. If that happens, we'll likely add a
// CoreMLProviderOptions struct and an AppendExecutionProviderCoreMLV2 function
// to the Go wrapper library, but for now the simpler API is the only thing
// available.
//
// Regardless, the meanings of the flag bits are currently defined in the
// coreml_provider_factory.h file which is provided in the include/ directory of
// the onnxruntime releases for Apple platforms.
func (o *SessionOptions) AppendExecutionProviderCoreML(flags uint32) error {
	status := C.AppendExecutionProviderCoreML(o.o, C.uint32_t(flags))
	if status != nil {
		return statusToError(status)
	}
	return nil
}
